#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A股股票智能筛选系统 - 东方财富API版本
直接调用东方财富公开API，增加重试机制和错误处理
"""

import requests
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import time
import json
from typing import Dict, List, Optional
import warnings
warnings.filterwarnings('ignore')


class EastMoneyAPI:
    """东方财富API封装"""
    
    def __init__(self):
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
            'Referer': 'http://quote.eastmoney.com/center/gridlist.html',
            'Accept': 'application/json, text/plain, */*',
            'Accept-Encoding': 'gzip, deflate',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Connection': 'keep-alive'
        })
        # 增加超时和重试
        self.retry_times = 5
        self.retry_delay = 3
    
    def _request_with_retry(self, url: str, params: dict = None) -> Optional[dict]:
        """带重试的请求"""
        for i in range(self.retry_times):
            try:
                response = self.session.get(url, params=params, timeout=15)
                if response.status_code == 200:
                    return response.json() if response.text else None
                else:
                    print(f"    HTTP {response.status_code}: {response.text[:100] if response.text else '无响应内容'}")
            except requests.exceptions.Timeout:
                if i < self.retry_times - 1:
                    print(f"    超时，{self.retry_delay}秒后重试... ({i+1}/{self.retry_times})")
                    time.sleep(self.retry_delay)
                else:
                    print(f"    ✗ 请求超时")
            except requests.exceptions.RequestException as e:
                if i < self.retry_times - 1:
                    print(f"    网络错误，{self.retry_delay}秒后重试... ({i+1}/{self.retry_times})")
                    time.sleep(self.retry_delay)
                else:
                    print(f"    ✗ 请求失败: {type(e).__name__} - {e}")
            except Exception as e:
                print(f"    ✗ 未知错误: {type(e).__name__} - {e}")
                break
        return None
    
    def get_stock_list(self) -> List[Dict]:
        """获取A股股票列表（代码和名称）"""
        # 更新为最新的API地址
        url = "http://push2.eastmoney.com/api/qt/clist/get"
        params = {
            'pn': '1',
            'pz': '6000',
            'po': '1',
            'np': '1',
            'ut': 'bd1d9ddb04089700cf9c27f6f7426281',
            'fltt': '2',
            'invt': '2',
            'fid': 'f3',
            'fs': 'm:0+t:6,m:0+t:80,m:1+t:2,m:1+t:23',
            'fields': 'f12,f14'
        }
        
        print("→ 正在获取A股股票列表...")
        data = self._request_with_retry(url, params)
        
        if not data or 'data' not in data or not data['data']['diff']:
            print("  ✗ 获取失败")
            return []
        
        stocks = []
        for item in data['data']['diff']:
            code = item.get('f12', '')
            name = item.get('f14', '')
            
            # 基础过滤：排除ST、退市、北交所
            if 'ST' in name or '退' in name or 'N' in name:
                continue
            if code.startswith(('4', '8')):
                continue
            
            stocks.append({'代码': code, '名称': name})
        
        print(f"  ✓ 获取 {len(stocks)} 支股票代码\n")
        return stocks
    
    def get_stock_info(self, symbol: str) -> Dict:
        """获取个股详细信息"""
        url = "http://push2.eastmoney.com/api/qt/stock/get"
        params = {
            'secid': f"1.{symbol}" if symbol.startswith('6') else f"0.{symbol}",
            'fields': 'f57,f58,f84,f85,f127,f116,f60,f45,f46,f47,f48,f50,f107,f161,f162'
        }
        
        data = self._request_with_retry(url, params)
        if not data or 'data' not in data:
            return {}
        
        info = data['data']
        return {
            '行业': info.get('f127', '未知'),
            '上市时间': str(info.get('f26', '')),
        }
    
    def get_stock_kline_data(self, symbol: str, name: str) -> Dict:
        """获取股票K线数据并计算所有指标"""
        url = "http://push2his.eastmoney.com/api/qt/stock/kline/get"
        
        # 判断市场
        secid = f"1.{symbol}" if symbol.startswith('6') else f"0.{symbol}"
        
        params = {
            'secid': secid,
            'klt': '101',  # 日K
            'fqt': '1',    # 前复权
            'lmt': '800',  # 获取800个交易日（约3年多数据）
            'fields1': 'f1,f2,f3,f4,f5,f6',
            'fields2': 'f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61'
        }
        
        data = self._request_with_retry(url, params)
        if not data or 'data' not in data or not data['data']:
            return None
        
        klines = data['data'].get('klines', [])
        if not klines or len(klines) < 250:  # 至少需要250个交易日（约1年）
            return None
        
        # 解析K线数据
        df_list = []
        for kline in klines:
            parts = kline.split(',')
            if len(parts) >= 7:
                df_list.append({
                    '日期': parts[0],
                    '开盘': float(parts[1]),
                    '收盘': float(parts[2]),
                    '最高': float(parts[3]),
                    '最低': float(parts[4]),
                    '成交量': float(parts[5]),
                    '成交额': float(parts[6]),
                    '振幅': float(parts[7]) if len(parts) > 7 else 0,
                    '涨跌幅': float(parts[8]) if len(parts) > 8 else 0,
                    '涨跌额': float(parts[9]) if len(parts) > 9 else 0,
                    '换手率': float(parts[10]) if len(parts) > 10 else 0,
                })
        
        df = pd.DataFrame(df_list)
        
        # 计算各项指标
        try:
            latest = df.iloc[-1]  # 最新一天
            
            # 基础数据
            result = {
                '代码': symbol,
                '名称': name,
                '最新价': latest['收盘'],
                '涨跌幅': latest['涨跌幅'],
                '成交额': latest['成交额'],
                '换手率': latest['换手率'],
                '振幅': latest['振幅'],
            }
            
            # 计算技术指标
            # MA60
            if len(df) >= 60:
                result['MA60'] = df['收盘'].tail(60).mean()
            else:
                result['MA60'] = df['收盘'].mean()
            
            # 52周最高价（约250个交易日）
            if len(df) >= 250:
                result['52周最高'] = df['收盘'].tail(250).max()
            else:
                result['52周最高'] = df['收盘'].max()
            
            # 3个月涨跌幅（约60个交易日）
            if len(df) >= 60:
                price_60d = df.iloc[-60]['收盘']
                result['3个月涨跌幅'] = ((result['最新价'] - price_60d) / price_60d) * 100
            else:
                result['3个月涨跌幅'] = 0
            
            # 52周回撤
            result['52周回撤'] = ((result['52周最高'] - result['最新价']) / result['52周最高']) * 100
            
            # 平均成交额（近20日）
            if len(df) >= 20:
                result['平均成交额'] = df['成交额'].tail(20).mean()
            else:
                result['平均成交额'] = df['成交额'].mean()
            
            # 平均换手率（近20日）
            if len(df) >= 20:
                result['平均换手率'] = df['换手率'].tail(20).mean()
            else:
                result['平均换手率'] = df['换手率'].mean()
            
            # 估算总市值（用成交额/换手率*100）
            if result['换手率'] > 0:
                result['总市值'] = (result['成交额'] / result['换手率']) * 100
            else:
                result['总市值'] = 0
            
            return result
            
        except Exception as e:
            return None


class ScreenerConfig:
    """筛选配置"""
    # 基础条件 - 已优化调整
    MIN_PRICE = 7
    MAX_PRICE = 50
    MIN_VOLUME = 30_000_000  # 3000万元（降低）
    MIN_LIST_YEARS = 1.5  # 1.5年（放宽）
    MIN_MARKET_CAP = 20_00_000_000  # 20亿元（大幅降低）
    
    # 基本面指标（如果后续添加）
    MIN_PE = 5
    MAX_PE = 35  # 稍微放宽
    MIN_PB = 0.8
    MAX_PB = 6  # 稍微放宽
    
    # 技术面指标
    MIN_3M_RETURN = -15  # 百分比（放宽）
    MAX_3M_RETURN = 25  # 百分比（放宽）
    MAX_52W_DRAWDOWN = 40  # 百分比（放宽）
    MIN_TURNOVER = 0.3  # 最小换手率%（降低）
    
    # 优先行业
    PRIORITY_INDUSTRIES = [
        '半导体', '集成电路', '芯片', '电子', '软件', '云计算', '人工智能', '计算机', '通信',
        '医药', '生物', 'CXO', '医疗', '器械',
        '锂电', '光伏', '储能', '新能源', '电池', '风电', '电力设备',
        '机械', '自动化', '军工', '航空', '航天',
        '消费', '食品', '饮料', '白酒', '家电', '纺织', '零售'
    ]
    
    OUTPUT_COUNT = 20


class StockScreener:
    """股票筛选器"""
    
    def __init__(self):
        self.api = EastMoneyAPI()
        self.config = ScreenerConfig()
        self.fetch_time = datetime.now()
        self.stats = {
            '初始股票池': 0,
            '基础条件过滤': 0,
            '行业筛选': 0,
            '技术面筛选': 0,
            '最终入选': 0
        }
    
    def run(self) -> pd.DataFrame:
        """执行筛选（基于历史K线数据）"""
        print("="*80)
        print(f"A股智能筛选系统 - 基于历史K线数据".center(80))
        print("="*80)
        print(f"数据获取时间: {self.fetch_time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"说明: 获取2-3年历史K线数据进行分析\n")
        
        # 1. 获取股票列表
        stock_list = self.api.get_stock_list()
        if not stock_list:
            print("✗ 无法获取股票列表")
            return pd.DataFrame()
        
        self.stats['初始股票池'] = len(stock_list)
        
        # 2. 获取K线数据并计算指标
        print(f"【阶段1】获取历史K线数据并计算指标...")
        print(f"  处理范围: 前100支股票（可修改）")
        
        all_data = []
        process_count = min(len(stock_list), 100)  # 限制处理数量
        
        for idx, stock in enumerate(stock_list[:process_count], 1):
            if idx % 10 == 0:
                print(f"  进度: {idx}/{process_count}, 已获取有效数据: {len(all_data)}")
            
            # 获取K线数据
            kline_data = self.api.get_stock_kline_data(stock['代码'], stock['名称'])
            if kline_data:
                # 获取行业信息
                info = self.api.get_stock_info(stock['代码'])
                kline_data['行业'] = info.get('行业', '未知')
                kline_data['上市时间'] = info.get('上市时间', '')
                
                all_data.append(kline_data)
            
            time.sleep(0.1)  # 避免请求过快
        
        if not all_data:
            print("  ✗ 未获取到有效数据")
            return pd.DataFrame()
        
        df = pd.DataFrame(all_data)
        print(f"  ✓ 成功获取 {len(df)} 支股票数据\n")
        
        # 3. 基础条件筛选
        print("【阶段2】基础条件筛选...")
        print(f"  市值要求: ≥{self.config.MIN_MARKET_CAP/100000000:.0f}亿, 成交额: ≥{self.config.MIN_VOLUME/10000:.0f}万")
        df = self._filter_basic_kline(df)
        self.stats['基础条件过滤'] = len(df)
        print(f"  ✓ 通过基础条件: {len(df)} 支\n")
        
        if df.empty:
            return df
        
        # 4. 行业筛选
        print("【阶段3】行业筛选...")
        df = self._filter_industry_kline(df)
        self.stats['行业筛选'] = len(df)
        print(f"  ✓ 优先行业: {len(df)} 支\n")
        
        if df.empty:
            return df
        
        # 5. 技术面筛选
        print("【阶段4】技术面筛选...")
        df = self._filter_technical_kline(df)
        self.stats['技术面筛选'] = len(df)
        print(f"  ✓ 通过技术面: {len(df)} 支\n")
        
        # 6. 综合评分排序
        print("【阶段5】综合评分与排序...")
        df = self._score_and_rank(df)
        df = df.head(self.config.OUTPUT_COUNT)
        self.stats['最终入选'] = len(df)
        print(f"  ✓ 最终入选: {len(df)} 支\n")
        
        return df
    
    def _filter_basic_kline(self, df: pd.DataFrame) -> pd.DataFrame:
        """基于K线数据的基础条件过滤"""
        # 股价区间
        df = df[(df['最新价'] >= self.config.MIN_PRICE) & 
                (df['最新价'] <= self.config.MAX_PRICE)]
        
        # 平均成交额
        df = df[df['平均成交额'] >= self.config.MIN_VOLUME]
        
        # 总市值（估算值）
        df = df[df['总市值'] >= self.config.MIN_MARKET_CAP]
        
        # 平均换手率
        df = df[df['平均换手率'] >= self.config.MIN_TURNOVER]
        
        return df.reset_index(drop=True)
    
    def _filter_industry_kline(self, df: pd.DataFrame) -> pd.DataFrame:
        """行业筛选（K线版本）"""
        def is_priority(industry: str) -> bool:
            if pd.isna(industry) or industry == '未知':
                return False
            return any(kw in industry for kw in self.config.PRIORITY_INDUSTRIES)
        
        df = df[df['行业'].apply(is_priority)]
        return df.reset_index(drop=True)
    
    def _filter_technical_kline(self, df: pd.DataFrame) -> pd.DataFrame:
        """技术面筛选（K线版本）"""
        # 3个月涨跌幅
        df = df[(df['3个月涨跌幅'] >= self.config.MIN_3M_RETURN) & 
                (df['3个月涨跌幅'] <= self.config.MAX_3M_RETURN)]
        
        # 52周回撤
        df = df[df['52周回撤'] <= self.config.MAX_52W_DRAWDOWN]
        
        # MA60支撑
        df = df[df['最新价'] >= df['MA60']]
        
        return df.reset_index(drop=True)
    
    def _filter_industry(self, df: pd.DataFrame) -> pd.DataFrame:
        """行业筛选"""
        results = []
        total = min(len(df), 200)  # 限制查询数量
        
        print(f"  正在获取行业信息... (处理前{total}支)")
        for idx, row in df.head(total).iterrows():
            symbol = row['代码']
            
            # 获取行业和上市时间
            info = self.api.get_stock_info(symbol)
            industry = info.get('行业', '未知')
            list_date = info.get('上市时间', '')
            
            # 检查上市时间
            if list_date:
                try:
                    list_dt = datetime.strptime(str(list_date), '%Y%m%d')
                    years = (datetime.now() - list_dt).days / 365
                    if years < self.config.MIN_LIST_YEARS:
                        continue
                except:
                    pass
            
            # 检查行业
            if any(keyword in industry for keyword in self.config.PRIORITY_INDUSTRIES):
                row['行业'] = industry
                row['上市时间'] = list_date
                results.append(row.to_dict())
            
            if (idx + 1) % 50 == 0:
                print(f"    进度: {idx+1}/{total}, 已筛选: {len(results)}")
        
        return pd.DataFrame(results) if results else pd.DataFrame()
    
    def _filter_technical(self, df: pd.DataFrame) -> pd.DataFrame:
        """技术面筛选"""
        results = []
        total = min(len(df), 100)  # 限制数量
        
        print(f"  正在获取历史行情... (处理前{total}支)")
        for idx, row in df.head(total).iterrows():
            symbol = row['代码']
            
            # 获取历史数据
            hist = self.api.get_historical_data(symbol, 250)
            
            if not hist.empty and len(hist) >= 60:
                current_price = row['最新价']
                
                # 计算MA60
                ma60 = hist['收盘'].tail(60).mean()
                
                # 计算52周最高价
                max_52w = hist['收盘'].tail(250).max() if len(hist) >= 250 else hist['收盘'].max()
                
                # 计算3个月涨跌幅
                if len(hist) >= 60:
                    price_60d = hist.iloc[-60]['收盘']
                    return_3m = ((current_price - price_60d) / price_60d) * 100
                else:
                    return_3m = 0
                
                # 计算52周回撤
                drawdown = ((max_52w - current_price) / max_52w) * 100
                
                # 技术面过滤
                if (return_3m >= self.config.MIN_3M_RETURN and 
                    return_3m <= self.config.MAX_3M_RETURN and
                    drawdown <= self.config.MAX_52W_DRAWDOWN and
                    current_price >= ma60):
                    
                    row['MA60'] = ma60
                    row['52周最高'] = max_52w
                    row['3个月涨跌幅'] = return_3m
                    row['52周回撤'] = drawdown
                    results.append(row.to_dict())
            
            if (idx + 1) % 20 == 0:
                print(f"    进度: {idx+1}/{total}, 已筛选: {len(results)}")
        
        return pd.DataFrame(results) if results else pd.DataFrame()
    
    def _score_and_rank(self, df: pd.DataFrame) -> pd.DataFrame:
        """综合评分（基于K线数据）"""
        if df.empty:
            return df
        
        # 技术得分 50分
        # 涨幅适中：5-15%最佳
        return_score = 1 - abs(df['3个月涨跌幅'] - 10) / 20
        df['涨幅得分'] = return_score.clip(0, 1) * 25
        
        # 回撤越小越好
        drawdown_score = 1 - df['52周回撤'] / self.config.MAX_52W_DRAWDOWN
        df['回撤得分'] = drawdown_score.clip(0, 1) * 25
        
        # 流动性得分 30分
        volume_score = np.log10(df['平均成交额'] / 100_000_000).clip(0, 2) / 2
        df['成交额得分'] = volume_score * 15
        
        turnover_score = df['平均换手率'].clip(0, 10) / 10
        df['换手率得分'] = turnover_score * 15
        
        # 市值得分 20分（市值适中更好，50-300亿最佳）
        market_cap_billion = df['总市值'] / 100_000_000
        cap_score = 1 - abs(market_cap_billion - 150) / 200  # 150亿为中心
        df['市值得分'] = cap_score.clip(0, 1) * 20
        
        # 综合得分
        df['综合得分'] = (df['涨幅得分'] + df['回撤得分'] + 
                         df['成交额得分'] + df['换手率得分'] + df['市值得分'])
        
        # 推荐理由
        df['推荐理由'] = df.apply(self._generate_reason, axis=1)
        
        # 排序
        df = df.sort_values('综合得分', ascending=False)
        
        return df
    
    def _generate_reason(self, row) -> str:
        """生成推荐理由（基于K线数据）"""
        reasons = []
        
        # 行业
        if pd.notna(row.get('行业')) and row['行业'] != '未知':
            reasons.append(row['行业'])
        
        # 涨幅
        ret = row.get('3个月涨跌幅', 0)
        if 0 <= ret <= 10:
            reasons.append('稳健上涨')
        elif ret > 10:
            reasons.append('强势上涨')
        elif ret < 0:
            reasons.append('回调')
        
        # 回撤
        drawdown = row.get('52周回撤', 100)
        if drawdown < 10:
            reasons.append('近年度高点')
        elif drawdown < 20:
            reasons.append('相对高位')
        
        # 流动性
        turnover = row.get('平均换手率', 0)
        if turnover > 5:
            reasons.append('交易活跃')
        elif turnover > 3:
            reasons.append('流动性好')
        
        # 市值
        market_cap = row.get('总市值', 0) / 100_000_000
        if 80 <= market_cap <= 200:
            reasons.append('中盘股')
        elif market_cap > 200:
            reasons.append('大盘股')
        
        return '，'.join(reasons) if reasons else '技术面良好'
    
    def print_stats(self):
        """打印统计"""
        print("="*80)
        print("筛选统计".center(80))
        print("="*80)
        for key, value in self.stats.items():
            print(f"  {key}: {value:>6} 支")
        if self.stats['初始股票池'] > 0:
            rate = (self.stats['最终入选'] / self.stats['初始股票池']) * 100
            print(f"  筛选率: {rate:>6.2f}%")
        print("="*80 + "\n")


class ReportGenerator:
    """报告生成"""
    
    def __init__(self, df: pd.DataFrame, stats: dict, fetch_time: datetime):
        self.df = df
        self.stats = stats
        self.fetch_time = fetch_time
    
    def print_console(self):
        """控制台输出"""
        from tabulate import tabulate
        
        if self.df.empty:
            print("✗ 无筛选结果")
            return
        
        print("="*100)
        print(f"A股智能筛选结果 - Top {len(self.df)}".center(100))
        print("="*100 + "\n")
        
        # 准备显示数据
        cols = ['代码', '名称', '行业', '最新价', '总市值', 
                '3个月涨跌幅', '52周回撤', '平均换手率', '综合得分', '推荐理由']
        
        display_df = self.df[[c for c in cols if c in self.df.columns]].copy()
        
        # 格式化
        for col in display_df.columns:
            if col == '最新价':
                display_df[col] = display_df[col].apply(lambda x: f"{x:.2f}")
            elif col == '总市值':
                display_df[col] = display_df[col].apply(lambda x: f"{x/100000000:.1f}亿")
            elif col in ['3个月涨跌幅', '52周回撤', '平均换手率']:
                display_df[col] = display_df[col].apply(lambda x: f"{x:.1f}%")
            elif col == '综合得分':
                display_df[col] = display_df[col].apply(lambda x: f"{x:.1f}")
        
        print(tabulate(display_df, headers='keys', tablefmt='grid', showindex=False))
        
        # 风险提示
        print(f"\n{'='*100}")
        print("⚠️  风险提示".center(100))
        print("="*100)
        print("  1. 本筛选结果仅供参考，不构成投资建议")
        print("  2. 股市有风险，投资需谨慎")
        print("  3. 请结合自身研究与风险承受能力决策")
        print("="*100 + "\n")
    
    def export_excel(self, filename: str = None):
        """导出Excel"""
        if self.df.empty:
            return
        
        if not filename:
            filename = f"A股筛选_{self.fetch_time.strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        try:
            with pd.ExcelWriter(filename, engine='openpyxl') as writer:
                # 筛选结果
                cols = ['代码', '名称', '行业', '最新价', '总市值', '平均成交额',
                       '3个月涨跌幅', '52周回撤', '平均换手率', 'MA60', 
                       '综合得分', '推荐理由']
                out_df = self.df[[c for c in cols if c in self.df.columns]].copy()
                # 格式化市值和成交额
                if '总市值' in out_df.columns:
                    out_df['总市值(亿)'] = (out_df['总市值'] / 100000000).round(2)
                    out_df = out_df.drop('总市值', axis=1)
                if '平均成交额' in out_df.columns:
                    out_df['平均成交额(亿)'] = (out_df['平均成交额'] / 100000000).round(2)
                    out_df = out_df.drop('平均成交额', axis=1)
                out_df.to_excel(writer, sheet_name='筛选结果', index=False)
                
                # 完整数据
                self.df.to_excel(writer, sheet_name='详细数据', index=False)
                
                # 参数说明
                params = pd.DataFrame({
                    '参数': ['数据源', '历史数据', '股价', '3月涨跌', '52周回撤', '成交额', '市值'],
                    '标准': [
                        '东方财富历史K线',
                        '800个交易日(约3年)',
                        f"{ScreenerConfig.MIN_PRICE}-{ScreenerConfig.MAX_PRICE}元",
                        f"{ScreenerConfig.MIN_3M_RETURN}%~{ScreenerConfig.MAX_3M_RETURN}%",
                        f"<{ScreenerConfig.MAX_52W_DRAWDOWN}%",
                        f">{ScreenerConfig.MIN_VOLUME/100000000:.0f}亿",
                        f">{ScreenerConfig.MIN_MARKET_CAP/100000000:.0f}亿"
                    ]
                })
                params.to_excel(writer, sheet_name='筛选参数', index=False)
                
                # 统计
                stats_df = pd.DataFrame(list(self.stats.items()), columns=['项目', '数量'])
                stats_df.to_excel(writer, sheet_name='统计信息', index=False)
            
            print(f"✓ Excel已导出: {filename}\n")
        except Exception as e:
            print(f"✗ 导出失败: {e}\n")


def main():
    """主函数"""
    # 执行筛选
    screener = StockScreener()
    result = screener.run()
    screener.print_stats()
    
    if result.empty:
        print("✗ 未找到符合条件的股票，建议调整筛选参数")
        return
    
    # 生成报告
    reporter = ReportGenerator(result, screener.stats, screener.fetch_time)
    reporter.print_console()
    reporter.export_excel()
    
    print("✓ 筛选任务完成！\n")


if __name__ == "__main__":
    main()
